# -*- coding: cp936 -*-

import lxm
Vector2 = lxm.Vector2
Vector3 = lxm.Vector3
#ZERO = 0.00000001
import math
import random
rand = random.Random()
import numpy as np
import numpy.ma as ma
class NodeForce():
    def __init__(self, m, g, size):
        self.m, self.g, self.size = m, g, size
    def __str__(self):
        return "NodeForce(%f, %f, %s)" % (self.m, self.g, str(self.size))
    def __repr__(self): return self.__str__()

    @classmethod
    def genNpGMNMatrix(self, lonfs):
        import numpy as np
        m = np.array([i.m for i in lonfs])
        m_1 = m.reshape(-1, 1)
        mn = m * m_1
        g = np.array([i.g for i in lonfs])
        g_1 = m.reshape(-1, 1)
        return (g + g_1) * mn
        
class EdgeForce():
    def __init__(self, src, tgt, k, nl):
        self.src, self.tgt = src, tgt
        self.k, self.nl = k, nl
    def __str__(self):
        return "EdgeForce(%d, %d, %f, %f)" % (self.src, self.tgt, self.k, self.nl)
    
class Force():
    MAX_FORCE = 50 #100
    @classmethod
    def f_analysis(self, a, b, r, f):
        vf = self.mul(self.sub(a, b), f/r)
        return vf, self.neg(vf), f
    @classmethod
    def calForce(self, a, b, fun):
        r = self.calR(a, b)
        f = fun(r)
        return self.f_analysis(a, b, r, f)
    @classmethod
    def calStringForce(self, a, b, k, l):
        """
        >>> a, b = (3, 4), (0, 0)
        >>> Force2.calStringForce(a, b, 1, 5) == ((0.0, 0.0), (-0.0, -0.0), 0.0)
        True
        >>> Force2.calStringForce(a, b, 1, 3)[2] == -2.0
        True
        >>> Force2.calStringForce(a, b, 1, 3)[0:2]==Force2.calStringForce(b, a, 1, 3)[1::-1]
        True
        """
        while (a==b): a, b = self.getRandom(), self.getRandom()
        return self.calForce(a, b, lambda r: -k * (r-l))
    @classmethod
    def calRepulsion(self, a, b, ma=1.0, mb=1.0, g=1.0):
        """
        >>> c, d = (3, 4), (3, 4.0000000000001)
        >>> Force2.calRepulsion(c, d)[2] == 100
        True
        """
        while (a==b): a, b = self.getRandom(), self.getRandom()
        if self.within(a, b, 0.00001):
            return self.f_analysis(a, b, self.calR(a, b), self.MAX_FORCE)
        return self.calForce(a, b, lambda r: g * ma * mb / (r*r))
    @classmethod
    def calRepulsion_Overlap(self, a, b, size, ma=1.0, mb=1.0, g=1.0):
        while (a==b): a, b = self.getRandom(), self.getRandom()
        if self.withinM(a, b, size):
            return self.f_analysis(a, b, self.calR(a, b), self.MAX_FORCE)
        return self.calForce(a, b, lambda r: g * ma * mb / sqr(r-l))
    @classmethod
    def within(self, a, b, r):
        return self.calR(a, b) < r

    @classmethod
    def calSquare(self, X):
        a = np.array(X)
        b = a.reshape(-1, 1)
        return a-b
    @classmethod
    def calRepulsion_np(self, loc, MAX_FORCE, GMN):
        diff, rr = self.calRR_np(loc)
        mx = ma.masked_equal(rr, 0.0)
        mx.fill_value = 1.0 / MAX_FORCE
        mx = mx.filled()
        f = GMN * np.power(mx, -1)
        forces = [f*(v/mx) for v in diff]
        return zip(*[[np.sum(e[:i])+np.sum(e[i+1:]) for i, e in enumerate(f)] for f in forces])

import random
global_random = random.Random(1)
def randF(a, b):
    return a + global_random.random() * (b-a)
    
class Force2(Force, Vector2):
    @classmethod
    def withinM(self, a, b, size):
        if abs(a[0]-b[0])<size[0]: return True
        if abs(a[1]-b[1])<size[1]: return True
        return False
    @classmethod
    def calRR_np(self, vec):
        a, b = map(lambda e:np.array(e, dtype=np.float), zip(*vec))
        a_1, b_1 = a.reshape(-1, 1), b.reshape(-1, 1)
        #a_2, b_2 = a-a_1, b-b_1
        a_2, b_2 = a_1 - a, b_1 - b
        return (a_2, b_2), np.square(a_2)+np.square(b_2)

    @classmethod
    def random(self, a, b):
        return (randF(a, b), randF(a, b))
        
    
class Force3(Force, Vector3):
    @classmethod
    def withinM(self, a, b, size):
        if abs(a[0]-b[0])<size[0]: return True
        if abs(a[1]-b[1])<size[1]: return True
        if abs(a[2]-b[2])<size[2]: return True
        return False
    @classmethod
    def calRR_np(self, vec):
        a, b, c = map(lambda e:np.array(e, dtype=np.float), zip(*vec))
        a_1, b_1, c_1 = a.reshape(-1, 1), b.reshape(-1, 1)
        #a_2, b_2, c_1 = a-a_1, b-b_1, c-c_1
        a_2, b_2, c_1 = a_1-a, b_1-b, c_1-c
        return (a_2, b_2, c_2), np.square(a_2)+np.square(b_2)+np.square(c_2)
    @classmethod
    def random(self, a, b):
        return (randF(a, b), randF(a, b))
        
class InitLayout():
    def __init__(self, nodes, edges, ForceUtil=Force2):
        self.nodes = nodes;
        self.edges = edges;
        self.forceUtil = ForceUtil
        self.loc = self.forceUtil.genList(len(nodes))
    def runAlgor(self):
        y = self.forceUtil
        self.loc = [y.mul(y.getNum(i), 10) for i, l in enumerate(self.loc)]
        return self.loc


class InitLayout_Random():
    def __init__(self, nodes, edges, ForceUtil=Force2):
        self.nodes = nodes;
        self.edges = edges;
        self.forceUtil = ForceUtil
        self.loc = self.forceUtil.genList(len(nodes))
    def runAlgor(self):
        y = self.forceUtil
        self.loc = [y.random(0, 10000) for i, l in enumerate(self.loc)]
        return self.loc

t=0
import time
import copy
class FDL():
    def __init__(self, nodelist, edgelist, parameters, ForceUtil=Force2, initloc=None):        
        self.nodenum = len(nodelist)

        self.nodelist = nodelist
        self.edgelist = edgelist
        self.nodefs = [NodeForce(parameters.NodeM, parameters.RepulsionG, parameters.NodeSize)] * len(nodelist)
        self.edgefs = [EdgeForce(a, b, parameters.SpringK, parameters.SpringLen) for a, b in edgelist]

        self.p_static = copy.deepcopy(parameters)
        
        self.forceUtil = ForceUtil
        self.acc_force = ForceUtil.genList(self.nodenum)
        if initloc == None:
            self.loc = ForceUtil.genList(self.nodenum)
        else: self.loc = initloc
        self.gmn = NodeForce.genNpGMNMatrix(self.nodefs)

    def runAlgorWithGen(self, gens):
        for i in xrange(gens):
            self.runAlgorStepByStep()
        return self.loc
    
    def runAlgorStepByStep(self):
        self.acc_force = self.calForceFromLocation(self.loc)
        self.loc = self.calLocationFromForce()
        #print self.loc
        #center = self.forceUtil.center(self.loc)
        #self.loc = [self.forceUtil.sub(l, center) for l in self.loc]

    ### Basic elements  
    def calForceFromLocation(self, loc):
        global t
        springf = self.calSpringForces(loc,
                            self.edgefs, self.nodefs,
                            self.forceUtil)
        a = time.time()
        #repulf = self.calRepulsionForces(loc, self.nodefs, self.forceUtil)
        repulf = self.forceUtil.calRepulsion_np(loc, Force.MAX_FORCE, self.gmn)
        t += time.time()-a
        return [self.forceUtil.add(a, b) for a, b in zip(springf, repulf)]
        
    def calLocationFromForce(self):
        return self.integrator.integrate(self.loc, self.acc_force, self.nodefs)

    ### Force Calculation
    @classmethod
    def calSpringForces(self, loc, edgefs, nodefs, ForceUtil=Force2):
        u = ForceUtil
        acc_force = u.genList(len(loc))
        for e in edgefs:
            loc0, loc1 = loc[e.src], loc[e.tgt]
            nfa, nfb = nodefs[e.src], nodefs[e.tgt]
            f0, f1, f = u.calStringForce(loc0, loc1, e.k,
                                         u.norm(nfa.size)+u.norm(nfb.size)+e.nl)
            acc_force[e.src] = u.add(acc_force[e.src], f0)
            acc_force[e.tgt] = u.add(acc_force[e.tgt], f1)
        return acc_force

    @classmethod
    def calRepulsionForces(self, loc, nodefs, ForceUtil=Force2):
        u = ForceUtil
        acc_force = u.genList(len(loc))
        for id0, nf0 in enumerate(nodefs):
            for id1, nf1 in enumerate(nodefs[id0+1:], id0+1):
                """f0, f1, f = u.calRepulsion_Overlap(loc[id0], loc[id1],
                                                   size=u.add(nf0.size, nf1.size),
                                                   ma=nf0.m, mb=nf1.m)"""
                f0, f1, f = u.calRepulsion(loc[id0], loc[id1],
                                           ma=nf0.m, mb=nf1.m,
                                           g=(nf0.g+nf1.g)/2.0)
                acc_force[id0] = u.add(acc_force[id0], f0)
                acc_force[id1] = u.add(acc_force[id1], f1)
        return acc_force

if __name__ == "__main__":
    import doctest
    doctest.testmod()



